import tensorflow as tf
from sklearn.datasets import load_iris
import os
from PsliteOptimizer import PsliteOptimizer
import argparse


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--rank', type=int, required=True, dest='rank')
    return parser.parse_args()


def env():
    args = get_args()
    os.environ['DMLC_PS_ROOT_URI'] = '127.0.0.1'
    os.environ['DMLC_PS_ROOT_PORT'] = '8899'
    os.environ['DMLC_ROLE'] = 'worker'
    os.environ['HEAPPROFILE'] = './W%d' % args.rank
    os.environ['DMLC_NUM_SERVER'] = '1'
    os.environ['DMLC_NUM_WORKER'] = '2'


if __name__ == '__main__':
    env()

    X, y = load_iris(return_X_y=True)
    X_plh = tf.placeholder(dtype=tf.float32, shape=[None, 4], name='x_input')
    y_plh = tf.placeholder(dtype=tf.int32, shape=[None], name='y_input')

    h1 = tf.layers.dense(X_plh, 64)
    h2 = tf.layers.dense(h1, 32)
    logits = tf.layers.dense(h2, 3, activation=None)

    loss_op = tf.losses.sparse_softmax_cross_entropy(labels=y_plh, logits=logits)
    opt = tf.train.GradientDescentOptimizer(learning_rate=0.001)
    opt = PsliteOptimizer(opt)
    train_op = opt.minimize(loss_op)
    rank_op = opt._ops.ps_my_rank()

    sess = tf.InteractiveSession()
    sess.run(tf.initialize_all_variables())
    opt.ps_init(sess)
    for i in range(100):
        _, loss, rank = sess.run([train_op, loss_op, rank_op], feed_dict={X_plh: X, y_plh: y})
        print('step: %d, loss: %f, rank: %d' % (i, loss, rank))

    sess.close()
